///////////////////////////////////////////////////////////////////////////////////////////////////
/// \file correlator.h
///
/// \brief   header file for the correlator class
///
///////////////////////////////////////////////////////////////////////////////////////////////////

#ifndef PCPS_CORRELATOR_HEADER
#define PCPS_CORRELATOR_HEADER

#include <src/pcps.h>
#include <src/cluster.h>
#include <src/blas_lapack.h>

namespace pcps {

  ////////////////////////////////////////////////////////////////////////////////////
  /// \brief   class for a correlator
  ///
  ////////////////////////////////////////////////////////////////////////////////////

  template <class S> class Correlator {

    private:

      bool _fixed; ///< whether the correlator is fixed or can be optimized
      bool _ti;    ///< whether this is a translated image of another correlator
      int _configs_per_site; ///< the number of configurations available to a single site
      int _nelements; ///< the number of data elements in the correlator
      boost::shared_ptr< std::vector<int> > _site_vec; ///< sites defining the correlator
      boost::shared_ptr< std::vector<int> > _ti_sites; ///< if a translated image, the sites defining the image's parent
      boost::shared_ptr< std::vector<S> > _data;  ///< storage for the correlator's values

    public:

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   default constructor
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator() : _fixed(false), _ti(false), _configs_per_site(0), _nelements(0), _site_vec(), _ti_sites(), _data() {}

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   constructor from a single site
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator(const int site, const int configs_per_site, const S value, const bool fixed = false) {
        std::vector<int> sites(1);
        sites.at(0) = site;
        this->initialize(sites, configs_per_site, value, fixed, false, pcps::Correlator<S>());
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   constructor from a set of sites
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator(const std::set<int> & site_set, const int configs_per_site, const S value, const bool fixed = false) {
        std::vector<int> sites(site_set.begin(), site_set.end());
        this->initialize(sites, configs_per_site, value, fixed, false, pcps::Correlator<S>());
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   constructor from a vector of sites
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator(const std::vector<int> & sites, const int configs_per_site, const S value, const bool fixed = false) {
        this->initialize(sites, configs_per_site, value, fixed, false, pcps::Correlator<S>());
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   translationally invariant constructor
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator(const pcps::Correlator<S> & other, const std::vector<int> & sites) {
        this->initialize(sites, other.configs_per_site(), pcps::unity<S>(), other.fixed(), true, other);
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   copy constructor (can be either deep or shallow)
      ///
      ///   \param  other  the correlator to copy from
      ///   \param  deep   whether to perform a deep or shallow copy (default is shallow)
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator(const Correlator<S> &other, const bool deep = false) {
        if (deep)
          this->deep_copy(other);
        else
          this->shallow_copy(other);
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   assignment operator (shallow copy)
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator<S> & operator=(const Correlator<S> &other) {
        this->shallow_copy(other);
        return *this;
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   assignment operator (deep copy)
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      Correlator<S> & operator<<=(const Correlator<S> &other) {
        this->deep_copy(other);
        return *this;
      }

      bool fixed() const { return _fixed; } ///< return whether the correlator is fixed or can be optimized
      bool ti() const { return _ti; } ///< return whether this is a translated image of another correlator
      int configs_per_site() const { return _configs_per_site; }
      int nelements() const { return _nelements; }
      const std::vector<int> & site_vec() const { return *_site_vec; }
      const std::vector<int> & ti_sites() const { assert( _ti ); return *_ti_sites; }
      S * data_ptr() {
        assert(_data->size() > 0);
        return &(*_data)[0];
      }
      const S * data_ptr() const {
        assert(_data->size() > 0);
        return &(*_data)[0];
      }
      int get_cmpd(const int * const config) const {
        assert( this->site_vec().size() > 0 );
        int cmpd = 0;
        const int * const site = &this->site_vec()[0];
        const int nsites = this->site_vec().size();
        for (int i = 0; i < nsites; i++) {
          assert( config[site[i]] >= 0 );
          assert( config[site[i]] < _configs_per_site );
          cmpd = cmpd * _configs_per_site + config[site[i]];
        }
        assert(cmpd >= 0 && cmpd < _data->size());
        return cmpd;
      }
      int get_cmpd(const pcps::Cluster & cluster) const {
        assert( this->site_vec().size() > 0 );
        assert( cluster.ket_occs().size() > *std::max_element(this->site_vec().begin(), this->site_vec().end()) );
        return this->get_cmpd(&cluster.ket_occs()[0]);
      }
      S & operator[](const int cmpd_index) {
        assert(cmpd_index >= 0 && cmpd_index < _data->size());
        return (*_data)[cmpd_index];
      }
      const S & operator[](const int cmpd_index) const {
        assert(cmpd_index >= 0 && cmpd_index < _data->size());
        return (*_data)[cmpd_index];
      }
      S & operator[](const pcps::Cluster & cluster) {
        return (*_data)[this->get_cmpd(cluster)];
      }
      const S & operator[](const pcps::Cluster & cluster) const {
        return (*_data)[this->get_cmpd(cluster)];
      }
      S & operator[](const int * const config) {
        return (*_data)[this->get_cmpd(config)];
      }
      const S & operator[](const int * const config) const {
        return (*_data)[this->get_cmpd(config)];
      }
      void normalize() {
        if (_nelements == 0) return;
        S max_value = pcps::zero<S>();
        S * p = &(*_data)[0];
        S * end = p + _nelements;
        for ( ; p != end; p++)
          if (pcps::square_norm(*p) > pcps::square_norm(max_value))
            max_value = *p;
        assert( std::abs(max_value) > 1.0e-5 );
        pcps::xscal(_nelements, pcps::unity<S>() / max_value, &(*_data)[0], 1);
      }
      void dump_binary(std::ofstream & ofs) const;
      void read_binary(std::ifstream & ifs);
      void broadcast(const MPI::Comm & comm, const int myrank, const int root);

    private:

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   function to initialize a correlator
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      void initialize(const std::vector<int> & sites, const int configs_per_site, const S value, const bool fixed, const bool ti, 
                      const pcps::Correlator<S> & other) {

        assert(configs_per_site > 0);
        assert(sites.size() > 0);

        // record whether this correlator should be held fixed during optimization
        _fixed = fixed;

        // record whether this correlator is a placeholder for another translationally invariant correlator
        _ti = ti;

        // store the number of configurations per site
        _configs_per_site = configs_per_site;

        // compute the number of data elements
        _nelements = 1;
        for (int n = sites.size(); n > 0; n--)
          _nelements *= configs_per_site;

        // ensure the sites are distinct
        if (std::set<int>(sites.begin(), sites.end()).size() != sites.size())
          throw pcps::Exception("Error in pcps::Correlator::initialize.  The sites are not distinct.");

        // prepare the vector of sites
        { boost::shared_ptr< std::vector<int> > temp( new std::vector<int>(sites.begin(), sites.end()) ); _site_vec = temp; }

        // translationally invariant correlator
        if (ti) {

          // ensure that this correlator has the same number of sites as its translationally invariang partner
          if (_site_vec->size() != other._site_vec->size())
            throw pcps::Exception("Error in pcps::Correlator::initialize.  Translational invariance requires equal numbers of sites.");

          // point to other correlator's sites
          _ti_sites = other._site_vec;

          // point to the other correlator's data array
          _data = other._data;

        // regular correlators
        } else {

          // prepare the data array, initializing all elements to value
          { boost::shared_ptr< std::vector<S> > temp( new std::vector<S>(_nelements, value) ); _data = temp; }

        }

      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   makes this correlator a shallow copy of another correlator
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      void shallow_copy(const pcps::Correlator<S> & other) {
        if (this == &other) return;
        _fixed = other._fixed;
        _ti = other._ti;
        _configs_per_site = other._configs_per_site;
        _nelements = other._nelements;
        _site_vec = other._site_vec;
        _ti_sites = other._ti_sites;
        _data = other._data;
      }

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   makes this correlator a deep copy of another correlator
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      void deep_copy(const pcps::Correlator<S> & other) {
        if (this == &other) return;
        _fixed = other._fixed;
        _ti = other._ti;
        _configs_per_site = other._configs_per_site;
        _nelements = other._nelements;
        _ti_sites = other._ti_sites;
        if ( other._site_vec.get() == 0 ) throw pcps::Exception("cannot deep copy a correlator with an empty _site_vec pointer");
        if (     other._data.get() == 0 ) throw pcps::Exception("cannot deep copy a correlator with an empty _data pointer"    );
        { boost::shared_ptr< std::vector<int   > > temp( new std::vector<int   >(*other._site_vec) ); _site_vec = temp; }
        if (other._ti)
          _data = other._data;
        else 
          { boost::shared_ptr< std::vector<S> > temp( new std::vector<S>(*other._data    ) ); _data     = temp; }
      }

    public:

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   function object for comparing correlators
      ///
      ////////////////////////////////////////////////////////////////////////////////////

      struct Compare_Obj {

        ////////////////////////////////////////////////////////////////////////////////////
        /// \brief   comparison operator defining "less than" for correlators
        ///
        ///   \param  x  1st correlator to compare
        ///   \param  y  2nd correlator to compare
        ///
        ///   \return whether x is less than y
        ///
        ////////////////////////////////////////////////////////////////////////////////////
        bool operator() (const Correlator<S> & x, const Correlator<S> & y) const {

          // first we place fixed correlators before non-fixed correlators
          if (x.fixed() && !y.fixed()) return true;
          if (y.fixed() && !x.fixed()) return false;

          // next we place translated images before regular correlators
          if (x.ti() && !y.ti()) return true;
          if (y.ti() && !x.ti()) return false;

          // next we place the correlators in ascending order based on the number of sites they have
          if (x.site_vec().size() < y.site_vec().size()) return true;
          if (y.site_vec().size() < x.site_vec().size()) return false;

          // finally, correlators are ordered alphabetically by their sorted site lists
          std::set<int> x_sites(x.site_vec().begin(), x.site_vec().end());
          std::set<int> y_sites(y.site_vec().begin(), y.site_vec().end());
          return std::lexicographical_compare(x_sites.begin(), x_sites.end(), y_sites.begin(), y_sites.end());

        }

      };

      ////////////////////////////////////////////////////////////////////////////////////
      /// \brief   a set that sorts correlators by their site indices
      ///
      ////////////////////////////////////////////////////////////////////////////////////
      typedef std::set< pcps::Correlator<S>, typename pcps::Correlator<S>::Compare_Obj > Set;

    // comparison operators
    public:

      bool operator==(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return ( !comp(*this, other) && !comp(other, *this) );
      }

      bool operator!=(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return ( comp(*this, other) || comp(other, *this) );
      }

      bool operator<(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return comp(*this, other);
      }

      bool operator>(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return comp(other, *this);
      }

      bool operator<=(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return !comp(other, *this);
      }

      bool operator>=(const pcps::Correlator<S> & other) const {
        typename pcps::Correlator<S>::Compare_Obj comp;
        return !comp(*this, other);
      }

  };

  template <class S>
  void absorb_correlators(const pcps::Input & userinp, std::vector< pcps::Correlator<S> > & correlators);

  template <class S>
  pcps::Correlator<S> spin_coherent_correlator(const int site,
                                                    const int configs_per_site,
                                                    const double theta,
                                                    const double phi,
                                                    const bool fixed);

  template <class S>
  void print_correlator(const pcps::Correlator<S> & corr);

}

#endif
